import time
# from IPython.display import display, Markdown
# from utils.tools import Tools
from assessment_assistant import LangchainQA, custom_rag_chain

def test_rag(qa, query):
    """测试RAG系统并返回结果"""
    start_time = time.time()
    ans = qa({"question": query})
    end_time = time.time()
    total_time = f"{round(end_time - start_time, 3)} sec."
    
    # 提取答案和源文档
    response = ans['answer']
    source_docs = ans.get('source_documents', [])
    
    # 构建完整响应
    full_response = f"Question: {query}\nAnswer: {response}\nTotal time: {total_time}"
    
    # 添加源文档信息（如果有）
    if source_docs:
        full_response += "\n\n来源文档:\n"
        for i, doc in enumerate(source_docs):
            source_info = f"{doc.metadata.get('source', '未知')}, 页码: {doc.metadata.get('page', '未知')}"
            full_response += f"{i+1}. {doc.page_content[:100]}... [{source_info}]\n"
    
    # 显示结果
    # tools = Tools()
    # display(Markdown(tools.colorize_text(full_response)))
    # print(full_response)
    return full_response


#优化查询后
def test_rag_opt(query):
   #new_langchain = custom_rag_chain()
   # new_langchain = LangchainQA()
    # qa = new_langchain.create_chain()
    start_time = time.time()
    #ans = new_langchain.get_result(query)
    ans = qa({"query": query})
    end_time = time.time()
    total_time = f"{round(end_time - start_time, 3)} sec."
    # 提取答案和源文档
    response = ans['result']
    source_docs = ans.get('source_documents', [])         
    # 构建完整响应
    full_response = f"Question: {query}\nAnswer: {response}\nTotal time: {total_time}"
    return full_response
 

# 简化的交互式测试函数，不使用select模块
def interactive_rag_test(qa):  #如果是测试test_rag_opt，qa是您已经初始化好的RetrievalQA实例
    """交互式RAG测试主函数"""
    print("RAG测试已启动。输入问题并按回车键获取答案，输入 'q' 并按回车键退出。")
    
    while True:
        try:
            # 获取用户输入
            query = input("\n请输入您的问题: ")
            
            # 检查是否要退出
            if query.lower() == 'q':
                break
                
            # 处理空输入
            if not query.strip():
                print("问题不能为空，请重新输入。")
                continue
                
            # 调用测试函数
            print(test_rag(qa, query))
            # print(test_rag_opt(query))
            
        except KeyboardInterrupt:
            print("\n程序被用户中断")
            break
        except Exception as e:
            print(f"发生错误: {str(e)}")
    
    print("RAG测试已结束。")

# 使用示例
langchain_qa = LangchainQA()
qa = langchain_qa.create_chain()
interactive_rag_test(qa)  # qa是您已经初始化好的RetrievalQA实例
